#include "wrapper_sstructmg.h"

const idx_t num_diag = 7;

int main(int argc, char* argv[])
{
    int my_pid, num_procs;
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);
    MPI_Comm_size(MPI_COMM_WORLD, &num_procs);

    idx_t arg_cnt = 1;
    const idx_t local_m = atoi(argv[arg_cnt++]),// 每个part内有一个box，该box有m*m*m个cell
                np_per_dim = atoi(argv[arg_cnt++]),// 每个part内有一个box，该box有p*p*p个进程组成
                glb_nparts = 3;
    const std::string   its_name = std::string(argv[arg_cnt++]),
                        prc_name = std::string(argv[arg_cnt++]);
    std::string config_mg_file;
	if (prc_name == "GMG") config_mg_file = std::string(argv[arg_cnt++]);

    const idx_t ncells_per_dim = local_m * np_per_dim,
                ncells_per_part = ncells_per_dim * ncells_per_dim * ncells_per_dim,
                nprocs_per_part = np_per_dim * np_per_dim * np_per_dim;
    assert(num_procs == nprocs_per_part * glb_nparts);

    const idx_t my_part = my_pid / nprocs_per_part,
                pid_in_part = my_pid - my_part * nprocs_per_part;
    idx_t cart_ids[3];
    {// 计算本进程在part内处于什么位置
        cart_ids[0] = pid_in_part / (np_per_dim * np_per_dim);
        cart_ids[1] = (pid_in_part - cart_ids[0] * np_per_dim * np_per_dim) / np_per_dim;
        cart_ids[2] = (pid_in_part - cart_ids[0] * np_per_dim * np_per_dim - cart_ids[1] * np_per_dim);
    }

    idx_t my_ilower[3], my_iupper[3], box_ends[3];
    idx_t my_nelems = 1;
    for (idx_t d = 0; d < 3; d++) {
        my_ilower[d] =  cart_ids[d]     * local_m;
        my_iupper[d] = (cart_ids[d] + 1)* local_m - 1;
        box_ends[d] = my_iupper[d] + 1;
        my_nelems *= (my_iupper[d] - my_ilower[d] + 1);
    }
#if LOAD_REBAL
    assert(np_per_dim < local_m);
    if (my_part == 0) {
        if (cart_ids[0] == np_per_dim - 1) { assert(my_iupper[0] == ncells_per_dim - 1);
            my_ilower[0] += np_per_dim - 1;
        } else {
            my_ilower[0] += cart_ids[0];
            my_iupper[0] = my_ilower[0] + local_m + 1 - 1;// 多计算1个
        }
        if (cart_ids[1] == np_per_dim - 1) { assert(my_iupper[1] == ncells_per_dim - 1);
            my_ilower[1] += np_per_dim - 1;
        } else {
            my_ilower[1] += cart_ids[1];
            my_iupper[1] = my_ilower[1] + local_m + 1 - 1;// 多计算1个
        }
    }
    else if (my_part == 1) {
        if (cart_ids[0] == np_per_dim - 1) { assert(my_iupper[0] == ncells_per_dim - 1);
            my_ilower[0] += np_per_dim - 1;
        } else {
            my_ilower[0] += cart_ids[0];
            my_iupper[0] = my_ilower[0] + local_m + 1 - 1;// 多计算1个
        }
        if (cart_ids[1] == 0) { assert(my_ilower[1] == 0);
            my_iupper[1] -= np_per_dim - 1;
        } else {
            my_ilower[1] -= (np_per_dim - cart_ids[1]);
            my_iupper[1] = my_ilower[1] + local_m + 1 - 1;// 多计算一个
        }
    }
    else if (my_part == 2) {
        if (cart_ids[0] == 0) { assert(my_ilower[0] == 0);
            my_iupper[0] -= np_per_dim - 1;
        } else {
            my_ilower[0] -= (np_per_dim - cart_ids[0]);
            my_iupper[0] = my_ilower[0] + local_m + 1 - 1;// 多计算一个
        }
        if (cart_ids[1] == np_per_dim - 1) { assert(my_iupper[1] == ncells_per_dim - 1);
            my_ilower[1] += np_per_dim - 1;
        } else {
            my_ilower[1] += cart_ids[1];
            my_iupper[1] = my_ilower[1] + local_m + 1 - 1;// 多计算1个
        }
    }
    // 重新计算
    my_nelems = 1;
    for (idx_t d = 0; d < 3; d++) {
        box_ends[d] = my_iupper[d] + 1;
        my_nelems *= (my_iupper[d] - my_ilower[d] + 1);
    }
#endif
#ifdef DEBUG_PRINT
    for (int p = 0; p < num_procs; p++) {
        if (my_pid == p) {
            printf("Proc %d part %d: (%d,%d,%d)~(%d,%d,%d)\n", my_pid, my_part,
                my_ilower[0], my_ilower[1], my_ilower[2],
                my_iupper[0], my_iupper[1], my_iupper[2] );
            fflush(stdout);
        }
        MPI_Barrier(MPI_COMM_WORLD);
    }
    MPI_Barrier(MPI_COMM_WORLD);
#endif
    if (my_pid == 0) {
        printf("\033[1;35m Local M = %d P: %d Total #procs: %d Total #dof: %d \033[0m \n",
            local_m, np_per_dim, num_procs, ncells_per_part * glb_nparts);
    }
    SStructMG_Init();
    {
        SStructGrid<idx_t> ssgrid(MPI_COMM_WORLD, 3, glb_nparts);
        const idx_t halos[3] = {1, 1, 1};
        {// 添加本进程所拥有的块和盒子
            const idx_t blk_exts[3] = { ncells_per_dim, ncells_per_dim, ncells_per_dim };
            const bool periods[3] = {false, false, false};
            ssgrid.addBlock(my_part, blk_exts, periods);
            ssgrid.addBox(my_part, my_ilower, box_ends);
        }
        ssgrid.assemble();
        const idx_t my_nblks = ssgrid.my_blk_list.size(); assert(my_nblks == 1);
        idx_t shalos [my_nblks * 3];
        for (idx_t k = 0; k < my_nblks; k++) {
            shalos[k * 3 + 0] = halos[0];
            shalos[k * 3 + 1] = halos[1];
            shalos[k * 3 + 2] = halos[2];
        }

        // 根据半结构网格建立向量
        par_SstructVector<idx_t, data_t> b(ssgrid, shalos, true);
        par_SstructVector<idx_t, data_t> x(b), y(b), x0(b);
        {// 填入向量数据
            data_t * vec_buf = new data_t [my_nelems];
            for (idx_t i = 0; i < my_nelems; i++)
                vec_buf[i] = 1.0;
            b.set_box_values(my_part, my_ilower, box_ends, vec_buf);
            x0.set_val(0.0, true);
            y.set_val(0.0, true);
            delete [] vec_buf; vec_buf = nullptr;
        }

        // 根据半结构网格建立矩阵
        __int128_t masks[my_nblks]; for (idx_t k = 0; k < my_nblks; k++) masks[k] = stencil_mask_3d7;
        par_SstructMatrix<idx_t, data_t, data_t> A(ssgrid, shalos, masks);
        int proc_diag_nrows = my_nelems;
        int proc_offd_nnz   = 0;
        {// 填入矩阵数据
            // 首先处理对角块，也即结构化部分
            data_t* mat_buf = new data_t [my_nelems * num_diag];
            #pragma omp parallel for collapse(2) schedule(static)
            for (idx_t i0 = my_ilower[0]; i0 <= my_iupper[0]; i0++)
            for (idx_t i1 = my_ilower[1]; i1 <= my_iupper[1]; i1++)
            for (idx_t i2 = my_ilower[2]; i2 <= my_iupper[2]; i2++) {
                data_t* dst_ptr = mat_buf + (((i0 - my_ilower[0]) * (my_iupper[1] - my_ilower[1] + 1)
                                            +  i1 - my_ilower[1] )* (my_iupper[2] - my_ilower[2] + 1)
                                            +  i2 - my_ilower[2]  ) * num_diag;
                dst_ptr[0] = dst_ptr[1] = dst_ptr[2] = -1.0;
                dst_ptr[3] =  6.0;
                dst_ptr[4] = dst_ptr[5] = dst_ptr[6] = -1.0;
                // 边界赋0
                if (i0 == 0                 ) dst_ptr[0] = 0.0;
                if (i0 == ncells_per_dim - 1) dst_ptr[6] = 0.0;
                if (i1 == 0                 ) dst_ptr[1] = 0.0;
                if (i1 == ncells_per_dim - 1) dst_ptr[5] = 0.0;
                if (i2 == 0                 ) dst_ptr[2] = 0.0;
                if (i2 == ncells_per_dim - 1) dst_ptr[4] = 0.0;
            }
            A.set_box_values(my_part, my_ilower, box_ends, mat_buf);
            delete [] mat_buf; mat_buf = nullptr;

            // 然后处理非对角快，也即非结构部分
            A.init_offd();
            
            std::vector< Entry<idx_t, data_t> > all_nnz;
            const data_t od_val = -1.0;
            if      (my_part == 0) {
                if (cart_ids[0] == np_per_dim - 1) {// my: iub
                    assert(my_iupper[0] + 1 == ncells_per_dim);
                    idx_t ngb_part = 2;// ngb: ilb+
                    idx_t bdr_ilower[3] = {my_iupper[0], my_ilower[1], my_ilower[2]};// 注意这里应该是在内部
                    idx_t bdr_iupper[3] = {my_iupper[0], my_iupper[1], my_iupper[2]};// 注意这里应该是在内部
                    idx_t ngb_ilower[3]  = {0, my_ilower[1], my_ilower[2] };
                    idx_t ngb_iupper[3]  = {0, my_iupper[1], my_iupper[2] };
                    idx_t index_maps[3] = {0, 1, 2};
                    idx_t index_dirs[3] = {1, 1, 1};
                    idx_t rev_maps[3];
                    for (idx_t d = 0; d < 3; d++) {
                        rev_maps[index_maps[d]] = d;
                    }// {1, 0, 2}
                    for (idx_t i0 = bdr_ilower[0]; i0 <= bdr_iupper[0]; i0++)
                    for (idx_t i1 = bdr_ilower[1]; i1 <= bdr_iupper[1]; i1++)
                    for (idx_t i2 = bdr_ilower[2]; i2 <= bdr_iupper[2]; i2++) {
                        StructDescriptor<idx_t> row, col;
                        row.id = my_part;
                        row.index[0] = i0; row.index[1] = i1; row.index[2] = i2;
                        col.id = ngb_part;
                        for (idx_t d = 0; d < 3; d++) {
                            col.index[d] = index_dirs[rev_maps[d]] > 0 ?
                                    (ngb_ilower[d] + row.index[rev_maps[d]] - bdr_ilower[rev_maps[d]])
                                :   (ngb_iupper[d] - row.index[rev_maps[d]] + bdr_ilower[rev_maps[d]]);
                        }
                        // printf("Proc %d add row %d (%d,%d,%d) -- col %d (%d,%d,%d) %.2e\n", my_pid,
                        //     row.id, row.index[0], row.index[1], row.index[2],
                        //     col.id, col.index[0], col.index[1], col.index[2], od_val );
                        all_nnz.push_back(Entry<idx_t, data_t>(row, col, od_val));
                    }// elem loop
                }
                if (cart_ids[1] == np_per_dim - 1) {// my: jub
                    assert(my_iupper[1] + 1 == ncells_per_dim);
                    idx_t ngb_part = 1;// ngb: jlb+
                    idx_t bdr_ilower[3] = {my_ilower[0], my_iupper[1], my_ilower[2]};// 注意这里应该是在内部
                    idx_t bdr_iupper[3] = {my_iupper[0], my_iupper[1], my_iupper[2]};// 注意这里应该是在内部
                    idx_t ngb_ilower[3]  = {my_ilower[0], 0, my_ilower[2]};
                    idx_t ngb_iupper[3]  = {my_iupper[0], 0, my_iupper[2]};
                    idx_t index_maps[3] = {0, 1, 2};
                    idx_t index_dirs[3] = {1, 1, 1};
                    idx_t rev_maps[3];
                    for (idx_t d = 0; d < 3; d++) {
                        rev_maps[index_maps[d]] = d;
                    }// {1, 0, 2}
                    for (idx_t i0 = bdr_ilower[0]; i0 <= bdr_iupper[0]; i0++)
                    for (idx_t i1 = bdr_ilower[1]; i1 <= bdr_iupper[1]; i1++)
                    for (idx_t i2 = bdr_ilower[2]; i2 <= bdr_iupper[2]; i2++) {
                        StructDescriptor<idx_t> row, col;
                        row.id = my_part;
                        row.index[0] = i0; row.index[1] = i1; row.index[2] = i2;
                        col.id = ngb_part;
                        for (idx_t d = 0; d < 3; d++) {
                            col.index[d] = index_dirs[rev_maps[d]] > 0 ?
                                    (ngb_ilower[d] + row.index[rev_maps[d]] - bdr_ilower[rev_maps[d]])
                                :   (ngb_iupper[d] - row.index[rev_maps[d]] + bdr_ilower[rev_maps[d]]);
                        }
                        // printf("Proc %d add row %d (%d,%d,%d) -- col %d (%d,%d,%d) %.2e\n", my_pid,
                        //     row.id, row.index[0], row.index[1], row.index[2],
                        //     col.id, col.index[0], col.index[1], col.index[2], od_val );
                        all_nnz.push_back(Entry<idx_t, data_t>(row, col, od_val));
                    }// elem loop
                }
            }
            else if (my_part == 1) {
                if (cart_ids[0] == np_per_dim - 1) {// my: iub
                    assert(my_iupper[0] + 1 == ncells_per_dim);
                    idx_t ngb_part = 2;// ngb: jub-
                    idx_t bdr_ilower[3] = {my_iupper[0], my_ilower[1], my_ilower[2]};// 注意这里应该是在内部
                    idx_t bdr_iupper[3] = {my_iupper[0], my_iupper[1], my_iupper[2]};// 注意这里应该是在内部
                    idx_t ngb_ilower[3]  = {my_ilower[1], ncells_per_dim - 1, my_ilower[2]};// 注意这里是邻居的内部
                    idx_t ngb_iupper[3]  = {my_iupper[1], ncells_per_dim - 1, my_iupper[2]};// 注意这里是邻居的内部
                    idx_t index_maps[3] = {1, 0, 2};
                    idx_t index_dirs[3] = {-1, 1, 1};
                    idx_t rev_maps[3];
                    for (idx_t d = 0; d < 3; d++) {
                        rev_maps[index_maps[d]] = d;
                    }// {1, 0, 2}
                    for (idx_t i0 = bdr_ilower[0]; i0 <= bdr_iupper[0]; i0++)
                    for (idx_t i1 = bdr_ilower[1]; i1 <= bdr_iupper[1]; i1++)
                    for (idx_t i2 = bdr_ilower[2]; i2 <= bdr_iupper[2]; i2++) {
                        StructDescriptor<idx_t> row, col;
                        row.id = my_part;
                        row.index[0] = i0; row.index[1] = i1; row.index[2] = i2;
                        col.id = ngb_part;
                        for (idx_t d = 0; d < 3; d++) {
                            col.index[d] = index_dirs[rev_maps[d]] > 0 ?
                                    (ngb_ilower[d] + row.index[rev_maps[d]] - bdr_ilower[rev_maps[d]])
                                :   (ngb_iupper[d] - row.index[rev_maps[d]] + bdr_ilower[rev_maps[d]]);
                        }
                        // printf("Proc %d add row %d (%d,%d,%d) -- col %d (%d,%d,%d) %.2e\n", my_pid,
                        //     row.id, row.index[0], row.index[1], row.index[2],
                        //     col.id, col.index[0], col.index[1], col.index[2], od_val );
                        all_nnz.push_back(Entry<idx_t, data_t>(row, col, od_val));
                    }// elem loop
                }
                if (cart_ids[1] == 0) {// my: jlb
                    assert(my_ilower[1] == 0);
                    idx_t ngb_part = 0;// ngb: jub+
                    idx_t bdr_ilower[3] = {my_ilower[0], 0, my_ilower[2]};// 注意这里应该是在内部
                    idx_t bdr_iupper[3] = {my_iupper[0], 0, my_iupper[2]};// 注意这里应该是在内部
                    idx_t ngb_ilower[3]  = {my_ilower[0], ncells_per_dim - 1, my_ilower[2]};
                    idx_t ngb_iupper[3]  = {my_iupper[0], ncells_per_dim - 1, my_iupper[2]};
                    idx_t index_maps[3] = {0, 1, 2};
                    idx_t index_dirs[3] = {1, 1, 1};
                    idx_t rev_maps[3];
                    for (idx_t d = 0; d < 3; d++) {
                        rev_maps[index_maps[d]] = d;
                    }// {1, 0, 2}
                    for (idx_t i0 = bdr_ilower[0]; i0 <= bdr_iupper[0]; i0++)
                    for (idx_t i1 = bdr_ilower[1]; i1 <= bdr_iupper[1]; i1++)
                    for (idx_t i2 = bdr_ilower[2]; i2 <= bdr_iupper[2]; i2++) {
                        StructDescriptor<idx_t> row, col;
                        row.id = my_part;
                        row.index[0] = i0; row.index[1] = i1; row.index[2] = i2;
                        col.id = ngb_part;
                        for (idx_t d = 0; d < 3; d++) {
                            col.index[d] = index_dirs[rev_maps[d]] > 0 ?
                                    (ngb_ilower[d] + row.index[rev_maps[d]] - bdr_ilower[rev_maps[d]])
                                :   (ngb_iupper[d] - row.index[rev_maps[d]] + bdr_ilower[rev_maps[d]]);
                        }
                        // printf("Proc %d add row %d (%d,%d,%d) -- col %d (%d,%d,%d) %.2e\n", my_pid,
                        //     row.id, row.index[0], row.index[1], row.index[2],
                        //     col.id, col.index[0], col.index[1], col.index[2], od_val );
                        all_nnz.push_back(Entry<idx_t, data_t>(row, col, od_val));
                    }// elem loop
                }
            }
            else if (my_part == 2) {
                if (cart_ids[0] == 0) {// my: ilb
                    assert(my_ilower[0] == 0);
                    idx_t ngb_part = 0;// ngb: iub+
                    idx_t bdr_ilower[3] = {0, my_ilower[1], my_ilower[2]};// 注意这里应该是在内部
                    idx_t bdr_iupper[3] = {0, my_iupper[1], my_iupper[2]};// 注意这里应该是在内部
                    idx_t ngb_ilower[3]  = {ncells_per_dim - 1, my_ilower[1], my_ilower[2]};// 注意这里是邻居的内部
                    idx_t ngb_iupper[3]  = {ncells_per_dim - 1, my_iupper[1], my_iupper[2]};// 注意这里是邻居的内部
                    idx_t index_maps[3] = {0, 1, 2};
                    idx_t index_dirs[3] = {1, 1, 1};
                    idx_t rev_maps[3];
                    for (idx_t d = 0; d < 3; d++) {
                        rev_maps[index_maps[d]] = d;
                    }// {1, 0, 2}
                    for (idx_t i0 = bdr_ilower[0]; i0 <= bdr_iupper[0]; i0++)
                    for (idx_t i1 = bdr_ilower[1]; i1 <= bdr_iupper[1]; i1++)
                    for (idx_t i2 = bdr_ilower[2]; i2 <= bdr_iupper[2]; i2++) {
                        StructDescriptor<idx_t> row, col;
                        row.id = my_part;
                        row.index[0] = i0; row.index[1] = i1; row.index[2] = i2;
                        col.id = ngb_part;
                        for (idx_t d = 0; d < 3; d++) {
                            col.index[d] = index_dirs[rev_maps[d]] > 0 ?
                                    (ngb_ilower[d] + row.index[rev_maps[d]] - bdr_ilower[rev_maps[d]])
                                :   (ngb_iupper[d] - row.index[rev_maps[d]] + bdr_ilower[rev_maps[d]]);
                        }
                        // printf("Proc %d add row %d (%d,%d,%d) -- col %d (%d,%d,%d) %.2e\n", my_pid,
                        //     row.id, row.index[0], row.index[1], row.index[2],
                        //     col.id, col.index[0], col.index[1], col.index[2], od_val );
                        all_nnz.push_back(Entry<idx_t, data_t>(row, col, od_val));
                    }// elem loop
                }
                if (cart_ids[1] == np_per_dim - 1) {// my: jub
                    assert(my_iupper[1] + 1 == ncells_per_dim);
                    idx_t ngb_part = 1;// ngb: iub-
                    idx_t bdr_ilower[3] = {my_ilower[0], my_iupper[1], my_ilower[2]};// 注意这里应该是在内部
                    idx_t bdr_iupper[3] = {my_iupper[0], my_iupper[1], my_iupper[2]};// 注意这里应该是在内部
                    idx_t ngb_ilower[3]  = {ncells_per_dim - 1, my_ilower[0], my_ilower[2]};// 注意这里是邻居的内部
                    idx_t ngb_iupper[3]  = {ncells_per_dim - 1, my_iupper[0], my_iupper[2]};// 注意这里是邻居的内部
                    idx_t index_maps[3] = {1, 0, 2};
                    idx_t index_dirs[3] = {1, -1, 1};
                    idx_t rev_maps[3];
                    for (idx_t d = 0; d < 3; d++) {
                        rev_maps[index_maps[d]] = d;
                    }// {1, 0, 2}
                    for (idx_t i0 = bdr_ilower[0]; i0 <= bdr_iupper[0]; i0++)
                    for (idx_t i1 = bdr_ilower[1]; i1 <= bdr_iupper[1]; i1++)
                    for (idx_t i2 = bdr_ilower[2]; i2 <= bdr_iupper[2]; i2++) {
                        StructDescriptor<idx_t> row, col;
                        row.id = my_part;
                        row.index[0] = i0; row.index[1] = i1; row.index[2] = i2;
                        col.id = ngb_part;
                        for (idx_t d = 0; d < 3; d++) {
                            col.index[d] = index_dirs[rev_maps[d]] > 0 ?
                                    (ngb_ilower[d] + row.index[rev_maps[d]] - bdr_ilower[rev_maps[d]])
                                :   (ngb_iupper[d] - row.index[rev_maps[d]] + bdr_ilower[rev_maps[d]]);
                        }
                        all_nnz.push_back(Entry<idx_t, data_t>(row, col, od_val));
                    }
                }
            }

            A.add_entries(all_nnz.size(), all_nnz.data());
            proc_offd_nnz = all_nnz.size();
        }
        A.assemble();

        // for (int p = 0; p < num_procs; p++) {
        //     if (p == my_pid) {
        //         printf("Proc %4d Diag nrows %8d Offd nnz %8d\n", p, proc_diag_nrows, proc_offd_nnz);
        //         fflush(stdout);
        //     }
        //     MPI_Barrier(MPI_COMM_WORLD);
        // }

        const int test_cnt = 10;
        std::vector<TEST_RECORD> records;
        for (int test = 0; test < test_cnt; test++) {
            vec_copy(x0, x);

            double ck = vec_dot<idx_t, data_t, double>(b, b);
            if (my_pid == 0) printf("(  b,   b) = %.15e\n", ck);

            ck = vec_dot<idx_t, data_t, double>(x, x);
            if (my_pid == 0) printf("(  x,   x) = %.15e\n", ck);

            A.Mult(b, y, false);
            ck = vec_dot<idx_t, data_t, double>(y, y);
            if (my_pid == 0) printf("(A*b, A*b) = %.15e\n", ck);

            A.Mult(x, y, false);
            ck = vec_dot<idx_t, data_t, double>(y, y);
            if (my_pid == 0) printf("(A*x, A*x) = %.15e\n", ck);

            TEST_CONFIG config;
            config.my_nblks = my_nblks;
            config.config_mg_file = config_mg_file;
            config.ssgrid_ptr = & ssgrid;
            config.rtol = 1.0e-7;
            TEST_RECORD rec;

            buildup_solver(its_name, prc_name, config);
            setup_and_solve(A, b, x, rec);

            data_t true_r_norm, b_norm;
            check_residual(A, x, b, y, true_r_norm, b_norm);

            if (my_pid == 0) {
                printf("\033[1;35mtrue ||r|| = %20.16e ||r||/||b||= %20.16e\033[0m\n", 
                    true_r_norm, true_r_norm / b_norm);
                printf("Proc %d Setup, Solve costs %.6f %.6f s\n",
                    my_pid, rec.setup, rec.solve);
            }

            stat_part_times(rec);
            records.push_back(rec);

            destroy_solver();
        }// test loop

        // 输出最佳时间
        if (my_pid == 0) {
            stat_multi_runs(records);
        }
    }// all Semi-** object must be destroyed before MPI_Finalize()
    SStructMG_Finalize();
    MPI_Finalize();
    return 0;
}